import torch
import torchvision.transforms as transforms
from PIL import Image
from torch import nn, tensor

from models.LeNet import LeNet
from models.MobileNet import MobileNetCIFAR  # 确保与训练代码的模型定义一致

# 设备配置（支持MPS/GPU/CPU）
device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")

# CIFAR-10类别标签（与训练时顺序一致）
classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')


# 加载训练时相同的预处理配置（网页7）
def create_transform():
    return transforms.Compose([
        transforms.Resize(32),  # 确保输入尺寸匹配训练尺寸
        transforms.CenterCrop(32),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.4914, 0.4822, 0.4465],  # 使用训练时的归一化参数
                             std=[0.247, 0.243, 0.261])
    ])


# 模型加载函数（网页3/5/7）
def load_model(model: nn.modules, model_path: str):
    model = model.to(device)
    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()  # 切换为评估模式
    return model


# 图像预测函数（网页3/5）
def predict_image(image_path, model, transform):
    try:
        # 加载并预处理图像（网页3）
        image = Image.open(image_path).convert('RGB')  # 强制转换为RGB格式
        input_tensor = transform(image).unsqueeze(0).to(device)  # 添加batch维度

        # 执行预测（网页5）
        with torch.no_grad():
            outputs = model(input_tensor)
            # print(classes)
            for i in range(0, len(classes)):
                print(f'{classes[i]} -> {outputs[0][i]}')
            _, predicted = torch.max(outputs, 1)

        # 获取预测结果（网页5）
        class_index = predicted.item()
        return classes[class_index]

    except Exception as e:
        print(f"预测过程中发生错误: {str(e)}")
        return None


# 主程序
if __name__ == "__main__":
    model = LeNet
    test_image_path = "dog.jpeg"

    # 初始化组件
    transform = create_transform()
    model = load_model(model(), f'../weights/cifar10_{model.__name__}.pth')  # 加载训练好的模型

    # 测试样例（可替换为任意图像路径）

    # 执行预测
    prediction = predict_image(test_image_path, model, transform)

    if prediction is not None:
        print(f"预测结果: {prediction}")
    else:
        print("预测失败，请检查输入图像路径或格式")
